#include <vector>
#include <stdio.h>
#include <string.h>
#include <GL/glew.h>
#include "Shader.h"


static const char vertexPrologue[] =
{
    "#version 330\n"
    "#define POSITION 0\n"
    "#define NORMAL 1\n"
    "#define TANGENT 2\n"
    "#define BITANGENT 3\n"
    "#define TEXCOORD0 4\n"
    "#define TEXCOORD1 5\n"
    "#define TEXCOORD2 6\n"
    "#define COLOUR0 7\n"
    "#define VERTEX 1\n"
    "#define FRAGMENT 0\n"
};

static const char fragmentPrologue[] =
{
    "#version 330\n"
    "#define VERTEX 0\n"
    "#define FRAGMENT 1\n"
};

typedef std::vector<char> stream;

static void Append(stream& s, const char* str, size_t len)
{
    s.insert(s.end(), str, str+len);
}

static void Append(stream& s, const char* str)
{
    Append(s, str, strlen(str));
}

static void Append(stream& s, char c)
{
    s.push_back(c);
}

static size_t CountNewlines(const char* begin, const char* end)
{
    size_t count = 0;
    for (const char* i=begin; i!=end; ++i)
    {
        if (*i == '\n')
            ++count;
    }
    return count;
}

static void CompilationError(stream& s, const char* path, const char* error)
{
    s.clear();
    Append(s, "Error compiling '");
    Append(s, path);
    Append(s, "': ");
    Append(s, error);
}


static bool Preprocess(stream& s, const char* path)
{
    char lineCountBuffer[64];
    stream buffer;

    FILE* file;
    if (fopen_s(&file, path, "rb") != 0)
    {
        CompilationError(s, path, "Failed to open file");
        return false;
    }

    fseek(file, 0, SEEK_END);
    size_t size = ftell(file);
    fseek(file, 0, SEEK_SET);

    buffer.resize(size+1);
    fread(&buffer[0], 1, size, file);
    fclose(file);
    buffer[size] = '\0';

    size_t lineCount = 1;
    char* begin = &buffer[0];
    char* include = strstr(begin, "#include");
    while (NULL != include)
    {
        lineCount += CountNewlines(begin, include);
        sprintf_s(lineCountBuffer, "#line %d\n", lineCount);
        Append(s, lineCountBuffer);

        *include = '\0';
        Append(s, begin);

        char* pathBegin = strchr(include+8, '\"');
        if (NULL == pathBegin)
        {
            CompilationError(s, path, "Malformed #include");
            return false;
        }
        ++pathBegin;

        char* pathEnd = strchr(pathBegin, '\"');
        if (NULL == pathEnd)
        {
            CompilationError(s, path, "Malformed #include");
            return false;
        }
        *pathEnd = '\0';

        if (!Preprocess(s, pathBegin))
            return false;

        begin = pathEnd + 1;
        include = strstr(begin, "#include");
    }
    sprintf_s(lineCountBuffer, "#line %d\n", lineCount);
    Append(s, lineCountBuffer);
    Append(s, begin);

    return true;
}


static char* CompileAndAttachShader(const char* path, GLuint program, const stream& fileStream, GLenum shaderType)
{
    stream s;
    if (shaderType == GL_VERTEX_SHADER)
        Append(s, vertexPrologue);
    else
        Append(s, fragmentPrologue);
    Append(s, &fileStream[0], fileStream.size());
    const char* source = &s[0];

    int param;
    GLuint shader = glCreateShader(shaderType);
    glShaderSource(shader, 1, &source, NULL);
    glCompileShader(shader);
    glGetShaderiv(shader, GL_COMPILE_STATUS, &param);

    char* error;
    if (param == GL_TRUE)
    {
        error = NULL;
        glAttachShader(program, shader);
    }
    else
    {
        glGetShaderiv(shader, GL_INFO_LOG_LENGTH, &param);
        char* info = new char[param];
        glGetShaderInfoLog(shader, param, NULL, info);
        CompilationError(s, path, info);
        Append(s, '\0');
        error = new char[s.size()];
        memcpy(error, &s[0], s.size());
        delete[] info;
    }

    glDeleteShader(shader);

    return error;
}


Shader* Shader::Compile(const char* path)
{
    stream fileStream;
    bool success = Preprocess(fileStream, path);
    Append(fileStream, '\n');
    Append(fileStream, '\0');

    Shader* shader = new Shader();
    if (!success)
    {
        shader->error = new char[fileStream.size()];
        memcpy(shader->error, &fileStream[0], fileStream.size());
    }
    else
    {
        int param;
        shader->error = CompileAndAttachShader(path, shader->handle, fileStream, GL_VERTEX_SHADER);
        if (!shader->IsValid())
            return shader;

        shader->error = CompileAndAttachShader(path, shader->handle, fileStream, GL_FRAGMENT_SHADER);
        if (!shader->IsValid())
            return shader;

        glLinkProgram(shader->handle);
        glGetProgramiv(shader->handle, GL_LINK_STATUS, &param);
        if (param == GL_FALSE)
        {
            stream s;
            glGetProgramiv(shader->handle, GL_INFO_LOG_LENGTH, &param);
            char* info = new char[param];
            glGetProgramInfoLog(shader->handle, param, NULL, shader->error);
            CompilationError(s, path, info);
            Append(s, '\0');
            shader->error = new char[s.size()];
            memcpy(shader->error, &s[0], s.size());
            delete[] info;
        }
    }

    return shader;
}


Shader::Shader()
    : error(NULL)
{
    handle = glCreateProgram();
}

Shader::~Shader()
{
    glDeleteProgram(handle);
    delete[] error;
}

bool Shader::IsValid() const
{
    return (error == NULL);
}

const char* Shader::GetError() const
{
    return error;
}


void BindShader(const Shader* shader)
{
    glUseProgram(shader->GetHandle());
}
